"""
文本优化服务

提供文本阅读优化功能，包括段落拆分和标题体系构建。
"""
import os
import yaml
from pathlib import Path
import ollama
from typing import Dict

class TextOptimizerService:
    """文本优化服务"""
    
    def __init__(self):
        """初始化文本优化服务"""
        # 加载配置
        config_path = Path(__file__).parent.parent.parent / "config" / "ollama.yaml"
        with open(config_path) as f:
            self.config = yaml.safe_load(f)
        
        # 设置Ollama环境变量
        os.environ["OLLAMA_HOST"] = f"{self.config['host']}:{self.config['port']}"
        
        # 模型配置
        self.model = self.config["model"]["name"]
        self.max_tokens = self.config["model"]["max_tokens"]
        self.temperature = self.config["model"]["temperature"]
    
    async def optimize_text(self, text: str) -> Dict:
        """优化文本，提升阅读性
        
        Args:
            text: 原始文本
            
        Returns:
            Dict: {
                'original': str,  # 原始文本
                'optimized': str  # 优化后的文本
            }
        """
        try:
            # 构建提示词
            prompt = f"""
            请针对原始文本，仅合理断句、添加标点符号和分段处理，提高可阅读性，输出应该同原始文本保持一致。


            原始文本：
            {text}
            """
            
            # 调用模型优化文本
            response = ollama.generate(
                model=self.model,
                prompt=prompt,
                options={
                    "num_ctx": self.max_tokens,
                    "num_predict": len(text)*1.3,
                    "temperature": self.temperature
                }
            )
            
            return {
                'original': text,
                'optimized': response['response']
            }
            
        except Exception as e:
            print(f"文本优化错误: {str(e)}")
            return {
                'original': text,
                'optimized': text  # 发生错误时返回原文
            } 